import torch
import cv2

def __pre_processing__(img, img_size, device):
    """
    对输入图像进行预处理，包括图像调整和填充、颜色空间转换和维度变换、转换为张量、归一化和添加批量维度。

    参数：
    img：输入图像，一个 HWC 格式的 NumPy 数组。
    img_size：模型要求的输入图像尺寸。
    device：张量计算的设备。

    返回值：
    经过预处理后的图像张量、缩放比例、左边界填充像素数和上边界填充像素数。
    """

    # 图像调整和填充
    img, r, left, top = __resize_with_padding__(img, (img_size, img_size))

    # 颜色空间转换和维度变换（BGR 转 RGB，HWC 转 CHW）
    img = img[:, :, ::-1].transpose((2, 0, 1)).copy()

    # 转换为张量，并将数据类型转换为 float
    img = torch.from_numpy(img).to(device).float()

    # 归一化（将像素值范围从 [0, 255] 缩放到 [0, 1]）
    img /= 255.0

    # 添加批量维度
    img = img.unsqueeze(0)

    return img, r, left, top

def __resize_with_padding__(img, size=(640,640)):
    """
    将图像调整为指定大小，并在调整后的图像周围填充灰色边框，保持原始图像的宽高比。

    参数：
    img：输入图像。
    size：目标大小，格式为 (width, height)。默认为 (640, 640)。

    返回值：
    调整后的图像、调整比例、左边界填充像素数、上边界填充像素数。
    """

    h, w, _ = img.shape  # 获取输入图像的高度和宽度
    r = min(size[0] / h, size[1] / w)  # 计算调整比例，选择最小的缩放比例，以便将整个图像放入目标尺寸内
    new_h, new_w = int(h * r), int(w * r)  # 计算调整后的图像大小
    new_img = cv2.resize(img, (new_w, new_h))  # 使用线性插值进行图像调整
    left = int((size[1] - new_w) / 2)  # 计算左边界填充像素数
    top = int((size[0] - new_h) / 2)  # 计算上边界填充像素数
    right = size[1] - left - new_w  # 计算右边界填充像素数
    bottom = size[0] - top - new_h  # 计算下边界填充像素数
    # 在调整后的图像周围填充灰色边框，保持图像的宽高比不变
    img = cv2.copyMakeBorder(new_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))

    return img, r, left, top  # 返回调整后的图像、调整比例、左边界填充像素数、上边界填充像素数

def __xywh2xyxy__(det):
    """
    将 `(x, y, w, h)` 格式的坐标转换为 `(x_min, y_min, x_max, y_max)` 格式的坐标。

    参数：
    det：包含一组目标框坐标的张量，其形状为 (N, 4)，N 表示目标框的数量，每行表示一个目标框，格式为 (x, y, w, h)。

    返回值：
    包含转换后目标框坐标的张量，其形状与输入相同，每行表示一个目标框，格式为 (x_min, y_min, x_max, y_max)。
    """
    y = det.clone()  # 创建副本以防止修改原始数据

    # 将 `(x, y, w, h)` 格式的坐标转换为 `(x_min, y_min, x_max, y_max)` 格式的坐标
    y[:, 0] = det[:, 0] - det[:, 2] / 2  # x_min = x - w / 2
    y[:, 1] = det[:, 1] - det[:, 3] / 2  # y_min = y - h / 2
    y[:, 2] = det[:, 0] + det[:, 2] / 2  # x_max = x + w / 2
    y[:, 3] = det[:, 1] + det[:, 3] / 2  # y_max = y + h / 2

    return y  # 返回转换后的目标框坐标张量

def __post_processing__(prediction, conf_threshold, iou_threshold, r, left, top, device):
    """
    对模型预测结果进行后处理，包括置信度阈值过滤、重新组织预测结果、非极大值抑制和坐标还原。

    参数：
    prediction：模型的预测结果张量，形状为 (N, 13, H, W)，N 表示批量大小，13 表示每个边界框的预测信息维度，H 和 W 表示特征图的高度和宽度。
    conf_threshold：置信度阈值，小于此阈值的边界框将被过滤掉。
    iou_threshold：IoU 阈值，用于进行非极大值抑制。
    r：缩放比例，用于将填充后的图像坐标还原为原始图像坐标。
    left：左边界填充像素数。
    top：上边界填充像素数。
    device：张量计算的设备。

    返回值：
    经过后处理的边界框预测结果。
    """

    # 将预测结果转置为 (N, H*W, 13) 的形状，并去除批量维度
    prediction = prediction.permute(0, 2, 1).squeeze(0)
    
    # 根据置信度阈值过滤掉置信度低于阈值的边界框
    xc = prediction[:, 4:6].amax(1) > conf_threshold
    x = prediction[xc]
    
    # 如果没有保留的边界框，则返回空列表
    if not len(x):
        return []

    # 提取边界框坐标，并将中心点宽高格式转换为左上右下格式
    boxes = x[:, :4]
    boxes = __xywh2xyxy__(boxes)
    
    # 提取边界框置信度、类别预测等信息，并重新组合成新的预测结果
    score, index = torch.max(x[:, 4:6], dim=-1, keepdim=True)
    x = torch.cat((boxes, score, x[:, 6:14], index), dim=1)

    # 对重新组织的预测结果进行非极大值抑制，去除重叠的边界框，保留具有最高置信度的边界框
    score = x[:, 4]
    keep = __nms__(x, iou_threshold, device)
    x = x[keep]

    # 将经过非极大值抑制后的边界框坐标从填充后的图像坐标还原到原始图像坐标
    x = __restore_box__(x, r, left, top)

    return x

def __restore_box__(dets, r, left, top):
    """
    将边界框的坐标从填充后的图像坐标还原到原始图像坐标。

    参数：
    dets：包含一组边界框坐标的张量，其形状为 (N, 4)，N 表示边界框的数量，每行表示一个边界框，格式为 (x_min, y_min, x_max, y_max)。
    r：缩放比例，用于将填充后的图像坐标还原为原始图像坐标。
    left：左边界填充像素数。
    top：上边界填充像素数。

    返回值：
    还原到原始图像坐标系的边界框坐标。
    """
    dets[:, [0, 2]] = dets[:, [0, 2]] - left  # 左边界填充像素数还原
    dets[:, [1, 3]] = dets[:, [1, 3]] - top  # 上边界填充像素数还原
    dets[:, :4] /= r  # 缩放比例还原

    return dets  # 返回还原到原始图像坐标系的边界框坐标

def __nms__(dets, iou_thresh, device):
    """
    非极大值抑制（NMS）操作，用于在目标检测中去除重叠的边界框。

    参数：
    dets：包含一组边界框及其得分的张量，其形状为 (N, 5)，N 表示边界框的数量，每行表示一个边界框，前四列为 (x_min, y_min, x_max, y_max)，最后一列为得分。
    iou_thresh：IoU 阈值，用于过滤重叠的边界框。
    device：张量计算的设备。

    返回值：
    保留的边界框索引列表。
    """

    y = dets.clone()  # 创建副本以防止修改原始数据
    y_box_score = y[:, :5]  # 获取边界框坐标和得分
    index = torch.argsort(y_box_score[:, -1], descending=True)  # 根据得分降序排列边界框索引

    keep = []  # 用于存储保留的边界框索引
    while index.size()[0] > 0:
        i = index[0].item()  # 获取得分最高的边界框索引
        keep.append(i)  # 将该索引加入保留列表

        # 计算当前边界框与其余边界框的交并比（IoU）
        x1 = torch.maximum(y_box_score[i, 0], y_box_score[index[1:], 0])
        y1 = torch.maximum(y_box_score[i, 1], y_box_score[index[1:], 1])
        x2 = torch.minimum(y_box_score[i, 2], y_box_score[index[1:], 2])
        y2 = torch.minimum(y_box_score[i, 3], y_box_score[index[1:], 3])

        zero_ = torch.tensor(0).to(device)
        w = torch.maximum(zero_, x2 - x1)
        h = torch.maximum(zero_, y2 - y1)
        inter_area = w * h

        nuion_area1 = (y_box_score[i, 2] - y_box_score[i, 0]) * (y_box_score[i, 3] - y_box_score[i, 1])  # 计算交集
        union_area2 = (y_box_score[index[1:], 2] - y_box_score[index[1:], 0]) * (
                    y_box_score[index[1:], 3] - y_box_score[index[1:], 1])  # 计算并集

        iou = inter_area / (nuion_area1 + union_area2 - inter_area)  # 计算IoU

        idx = torch.where(iou <= iou_thresh)[0]  # 保留IoU小于阈值的边界框
        index = index[idx + 1]  # 更新索引列表

    return keep  # 返回保留的边界框索引列表

def detect(img, detect_model, img_size, conf_threshold, iou_threshold, device):
    img, r, left, top = __pre_processing__(img, img_size, device)  # 前处理
    predict = detect_model(img)[0]  # 目标检测
    return __post_processing__(predict, conf_threshold, iou_threshold, r, left, top, device)  # 后处理